package com.aiserver.advice;

import com.aiserver.util.EncryptUtil;
import com.alibaba.fastjson.JSON;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import javax.servlet.http.HttpServletRequest;
import java.util.HashMap;
import java.util.Map;

/**
 * 响应体加密处理器
 * 自动对需要加密的响应进行加密处理
 */
@ControllerAdvice
public class EncryptResponseAdvice implements ResponseBodyAdvice<Object> {

    private static final Logger logger = LoggerFactory.getLogger(EncryptResponseAdvice.class);

    @Autowired
    private EncryptUtil encryptUtil;

    @Autowired
    private HttpServletRequest request;

    @Override
    public boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
        // 检查是否需要加密响应
        String encrypted = request.getHeader("X-Encrypted");
        return "true".equals(encrypted);
    }

    @Override
    public Object beforeBodyWrite(Object body, MethodParameter returnType, MediaType selectedContentType,
                                Class<? extends HttpMessageConverter<?>> selectedConverterType,
                                ServerHttpRequest request, ServerHttpResponse response) {
        
        try {
            // 如果响应体为空，直接返回
            if (body == null) {
                return createEncryptedResponse(null);
            }

            // 将响应对象转换为JSON字符串
            String responseJson;
            if (body instanceof String) {
                responseJson = (String) body;
            } else {
                responseJson = JSON.toJSONString(body);
            }

            // 加密响应数据
            String encryptedData = encryptUtil.aesEncrypt(responseJson);
            logger.debug("响应数据加密成功");

            // 返回加密后的响应格式
            return createEncryptedResponse(encryptedData);

        } catch (Exception e) {
            logger.error("响应数据加密失败", e);
            // 加密失败时返回原始数据
            return body;
        }
    }

    /**
     * 创建加密响应格式
     */
    private Map<String, Object> createEncryptedResponse(String encryptedData) {
        Map<String, Object> result = new HashMap<>();
        result.put("encrypted", true);
        result.put("data", encryptedData);
        result.put("timestamp", System.currentTimeMillis());
        return result;
    }
}